#!/usr/bin/env python

#$Id$
# -------------------------------------------------------------
#
__version__      = '$Revision$ '[11:-3]
__version_date__ = '$Date$ '[7:-3]
__author__       = 'R. Bouwens, <bouwens@ucolick.org>, D. Magee, <magee@ucolick.org>'


import pyfits
import imagestats
from pyraf import iraf
import os, sys, shutil, sampseq
from optparse import OptionParser

#Setup CALNICA header keyword switches
calnica_step_1 = ['BIASCORR', 'ZSIGCORR', 'ZOFFCORR', 'MASKCORR', 'NOISCALC', 
'NLINCORR', 'DARKCORR', 'BARSCORR', 'FLATCORR', 'UNITCORR', 'PHOTCALC']
calnica_step_2 = ['CRIDCALC', 'BACKCALC', 'WARNCALC']

MDIFF = 0.05 # Difference in median. If change in median > MDIFF run pedsky on individual multiaccum reads

def calPed(fitsimg, nrdir, pedsub):
    nrdir = os.path.join(os.path.abspath(nrdir), '')
    os.putenv('nref', nrdir)
    iraf.stsdas()
    iraf.stsdas.hst_calib()
    iraf.stsdas.hst_calib.nicmos()
    nicmos = iraf.stsdas.hst_calib.nicmos
    # Setup good regions for stats 
    nicmos.statregions(quad1='[1:127,15:127]',quad2='[129:256,15:127]', quad3='[1:127,129:256]',
                       quad4='[129:256,129:241]', statsec='[1:241,15:256]') 
    # Turn off crrejection for the first pass of calnica
    rawff = pyfits.open(fitsimg, 'update')
    phdu = rawff[0].header
    for kw in calnica_step_1:
        phdu.update(kw, 'PERFORM')
    for kw in calnica_step_2:
        phdu.update(kw, 'OMIT')
    rawff.flush()
    rawff.close()
    # First pass of calnica without cridcalc r
    print '**** Running CALNICA up to CRIDCALC step on %s...' % fitsimg
    nicmos.calnica(fitsimg, '')
    raw1 = fitsimg.replace('raw.fits', 'raw1.fits')
    shutil.copy(fitsimg, raw1)
    # Open IMA file and compare the median value of the first and last 3
    # exposures to determine if sky background changes in the exposure.
    ima = fitsimg.replace('raw.fits', 'ima.fits')
    ima0 = fitsimg.replace('raw.fits', 'ima0.fits')
    shutil.move(ima, ima0)
    imaff = pyfits.open(ima0, 'update')
    samp_seq = imaff[0].header['SAMP_SEQ']
    NBeginReads = sampseq.NBeginReads(samp_seq)
    exts = range(1, len(imaff)-1, 5)[:-NBeginReads]
    print '**** Comparing medians at beginning and end of exposure...'
    print 'sci extensions: ', exts
    BKG = imaff[1].header['GOODMEDN']
    for ext in exts:
        print imaff[ext].header['GOODMEDN'], imaff[ext].header['GOODSTDV']
    medn_top = 0
    for ext in exts[:2]:
        print 'top: ', ext
        medn_top += imaff[ext].header['GOODMEDN']
    medn_top = medn_top/2.0
    medn_bot = 0
    for ext in exts[-2:]:
        print 'bot: ', ext
        medn_bot += imaff[ext].header['GOODMEDN']
    medn_bot = medn_bot/2.0
    meddiff = abs(medn_bot - medn_top)
    print 'medn_top: %s medn_bot: %s medn_diff: %s' % (medn_top, medn_bot, meddiff)
    cal = fitsimg.replace('raw.fits', 'cal.fits')
    cal0 = fitsimg.replace('raw.fits', 'cal0.fits')
    os.remove(cal)
    if meddiff > MDIFF:
        print '******* %s sky deviation > %s *******' % (fitsimg, MDIFF)
        print '*******  Running pedsky on IMA  *******'
        # Sky background is changing so we neet to run pedsky on each multiaccum read
        # in the ima file then to flatten for cr rejection.
        for ev, ext in enumerate(exts): 
            hdulist = pyfits.HDUList([imaff[0], imaff[ext], imaff[ext+1], imaff[ext+2], imaff[ext+3], imaff[ext+4]])
            for i in hdulist[1:]:
                i.header['EXTVER'] = 1
            tmpin = 'temp_ext%s.fits' % ext
            tmpout = tmpin.replace('.fits', '_psky.fits')
            hdulist.writeto(tmpin)
            try:
                if pedsub:
                    nicmos.pedsub(tmpin, tmpout, filter='mask', eqorder=3)
                else:
                    nicmos.pedsky(tmpin, tmpout, salgorithm='quick')
            except:
              scipsky = pyfits.getdata(tmpin)
              imaff[ext].data = scipsky
              for i in range(ext, ext+5):
                  imaff[i].header['EXTVER'] = ev+1
            else:
              scipsky = pyfits.getdata(tmpout)
              imaff[ext].data = scipsky
              for i in range(ext, ext+5):
                  imaff[i].header['EXTVER'] = ev+1
            os.remove(tmpin)
            if os.path.exists(tmpout):
              os.remove(tmpout)
        imaff[0].header.update('CALPED', 1) 
        imaff[0].header.update('BKGND', BKG) 
        imaff.flush()
        imaff.close()
        # rename ima to raw
        shutil.move(ima0, fitsimg)
    else:
        # rename cal to raw
        shutil.move(ima0, fitsimg)
        calff = pyfits.open(fitsimg, 'update')
        calff[0].header.update('CALPED', 0)
        calff[0].header.update('BKGND', BKG) 
        calff.flush()
        calff.close()
    #Second calnica run for cr rejection
    rawff = pyfits.open(fitsimg, 'update')
    phdu = rawff[0].header
    for kw in calnica_step_2:
        phdu.update(kw, 'PERFORM')
    for kw in calnica_step_1:
        phdu.update(kw, 'OMIT')
    rawff.flush()
    rawff.close()
    print 'Running CALNICA from CRIDCALC step on %s' % fitsimg
    nicmos.calnica(fitsimg, '')
    calff = pyfits.open(cal, 'update')
    phdu = calff[0].header
    for kw in calnica_step_1:
        phdu.update(kw, 'PERFORM')
    calff.flush()
    calff.close()
    shutil.move(cal, cal0)
    if pedsub:
        nicmos.pedsub(cal0, cal, filter='mask', eqorder=3)
    else:
        nicmos.pedsky(cal0, cal, salgorithm='quick')
    # Clean up
    os.remove(fitsimg)
    os.remove(cal0)
    shutil.move(raw1, fitsimg)

if __name__ == '__main__':
    usage = 'usage: %prog fitsfile nrefdir'
    version = '%prog 0.1'
    parser = OptionParser(usage=usage, version=version)
    (options, args) = parser.parse_args()
    if len(args) != 2:
        parser.print_help()
        sys.exit()
    else:
        if not os.path.exists(args[0]):
            parser.error('file %s not found!' % (args[0]))
            sys.exit()
        if not os.path.exists(args[1]):
            parser.error('directory %s not found!' % (args[1]))
            sys.exit()
    fitsfile = args[0]
    nrefdir = args[1]
    calPed(fitsfile, nrefdir)
